import tkinter as tk
from tkinter import messagebox
import numpy as np
import random
import torch
import torch.nn as nn
import torch.optim as optim
import os
import time
from collections import deque
import threading

# 检查 GPU 是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 神经网络模型
class NeuralNetwork(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(NeuralNetwork, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

# 迷宫生成类
class MazeGenerator:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.maze = np.ones((height, width), dtype=int)

    def generateMaze(self):
        if self.width % 2 == 0:
            self.width += 1
        if self.height % 2 == 0:
            self.height += 1
        self.maze = np.ones((self.height, self.width), dtype=int)
        stack = []
        start_x, start_y = 1, 1
        self.maze[start_y][start_x] = 0
        stack.append((start_x, start_y))
        while stack:
            x, y = stack[-1]
            directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
            random.shuffle(directions)
            found = False
            for dx, dy in directions:
                nx, ny = x + 2 * dx, y + 2 * dy
                if 0 < nx < self.width - 1 and 0 < ny < self.height - 1 and self.maze[ny][nx] == 1:
                    self.maze[y + dy][x + dx] = 0
                    self.maze[ny][nx] = 0
                    stack.append((nx, ny))
                    found = True
                    break
            if not found:
                stack.pop()
        # 设置终点
        self.maze[self.height - 2][self.width - 2] = 2
        return self.maze

# 使用广度优先搜索（BFS）算法找到最短路径
def find_shortest_path(maze, start, end):
    queue = deque([(start, [start])])
    visited = set()
    while queue:
        (x, y), path = queue.popleft()
        if (x, y) == end:
            return path
        if (x, y) in visited:
            continue
        visited.add((x, y))
        for dx, dy in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
            new_x, new_y = x + dx, y + dy
            if 0 <= new_x < maze.shape[1] and 0 <= new_y < maze.shape[0] and maze[new_y][new_x] != 1:
                new_path = list(path)
                new_path.append((new_x, new_y))
                queue.append(((new_x, new_y), new_path))
    return None

# 迷宫游戏类
class MazeGame:
    def __init__(self, root, maze_index, num_mazes, show_all=False, first_game=None):
        self.root = root
        self.maze_index = maze_index
        self.num_mazes = num_mazes
        self.show_all = show_all
        self.first_game = first_game
        if show_all or maze_index == 0:
            self.canvas = tk.Canvas(root, width=300, height=300)
            self.canvas.pack(side=tk.LEFT, padx=5, pady=5)
        else:
            self.canvas = None

        if maze_index == 0:
            self.input_frame = tk.Frame(root)
            self.input_frame.pack()
            tk.Label(self.input_frame, text="长:").grid(row=0, column=0)
            self.width_entry = tk.Entry(self.input_frame)
            self.width_entry.grid(row=0, column=1)
            tk.Label(self.input_frame, text="宽:").grid(row=1, column=0)
            self.height_entry = tk.Entry(self.input_frame)
            self.height_entry.grid(row=1, column=1)
            tk.Button(self.input_frame, text="生成", command=self.generate_all_mazes).grid(row=2, column=0, columnspan=2)

            self.control_frame = tk.Frame(root)
            self.control_frame.pack()
            tk.Button(self.control_frame, text="上", command=lambda: self.move_player(0, -1)).grid(row=0, column=1)
            tk.Button(self.control_frame, text="下", command=lambda: self.move_player(0, 1)).grid(row=2, column=1)
            tk.Button(self.control_frame, text="左", command=lambda: self.move_player(-1, 0)).grid(row=1, column=0)
            tk.Button(self.control_frame, text="右", command=lambda: self.move_player(1, 0)).grid(row=1, column=2)
            self.ai_button = tk.Button(self.control_frame, text="启用神经网络", command=self.toggle_ai)
            self.ai_button.grid(row=3, column=0, columnspan=3)
            self.speed_scale = tk.Scale(self.control_frame, from_=1, to=10, orient=tk.HORIZONTAL, label="AI 速度", command=self.set_ai_speed)
            self.speed_scale.set(1)
            self.speed_scale.grid(row=4, column=0, columnspan=3)
        else:
            self.input_frame = None
            self.width_entry = None
            self.height_entry = None
            self.control_frame = None
            self.ai_button = None
            self.speed_scale = None

        self.score_label = tk.Label(root, text=f"AI 分数 (迷宫 {maze_index + 1}): 0")
        self.score_label.pack()

        self.ai_enabled = False
        self.ai_speed = 1
        self.ai_score = 0
        self.player_x = 1
        self.player_y = 1
        self.maze = None
        self.cell_size = 0
        self.model = NeuralNetwork(4, 64, 4).to(device)
        self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
        self.load_model()
        self.ai_steps = 0
        self.max_steps = 200
        self.start_time = None
        self.time_limit = 60
        self.epsilon = 0.1  # ε - 贪心策略参数
        self.shortest_path = None
        self.path_index = 0
        self.visited = set()
        self.recent_visited = deque(maxlen=5)  # 记录最近走过的 5 个位置
        self.ai_thread = None

    def generate_all_mazes(self):
        if self.maze_index == 0:
            try:
                width = int(self.width_entry.get())
                height = int(self.height_entry.get())
                for game in games:
                    game.generate_single_maze(width, height)
            except ValueError:
                messagebox.showerror("输入错误", "请输入有效的整数长和宽")

    def generate_single_maze(self, width, height):
        try:
            generator = MazeGenerator(width, height)
            self.maze = generator.generateMaze()
            self.cell_size = min(300 // self.maze.shape[1], 300 // self.maze.shape[0])
            if self.canvas:
                self.canvas.delete("all")
                self.draw_maze()
            self.player_x = 1
            self.player_y = 1
            if self.canvas:
                self.draw_player()
            self.ai_steps = 0
            self.start_time = time.time()
            # 生成迷宫后找到最短路径
            start = (1, 1)
            end = (self.maze.shape[1] - 2, self.maze.shape[0] - 2)
            self.shortest_path = find_shortest_path(self.maze, start, end)
            self.path_index = 0
            self.visited = set()
            self.recent_visited.clear()
        except Exception as e:
            print(f"生成迷宫时出错: {e}")

    def draw_maze(self):
        if self.canvas:
            for y in range(self.maze.shape[0]):
                for x in range(self.maze.shape[1]):
                    if self.maze[y][x] == 1:
                        color = "black"
                    elif self.maze[y][x] == 0:
                        color = "white"
                    else:
                        color = "red"
                    self.canvas.create_rectangle(x * self.cell_size, y * self.cell_size,
                                                 (x + 1) * self.cell_size, (y + 1) * self.cell_size,
                                                 fill=color)

    def draw_player(self):
        if self.canvas:
            try:
                self.canvas.delete("player")
                x0 = self.player_x * self.cell_size
                y0 = self.player_y * self.cell_size
                x1 = (self.player_x + 1) * self.cell_size
                y1 = (self.player_y + 1) * self.cell_size
                self.canvas.create_oval(x0, y0, x1, y1, fill="blue", tags="player")
            except tk.TclError:
                # 若 canvas 已被销毁，忽略错误
                pass

    def move_player(self, dx, dy):
        new_x = self.player_x + dx
        new_y = self.player_y + dy
        if 0 <= new_x < self.maze.shape[1] and 0 <= new_y < self.maze.shape[0] and self.maze[new_y][new_x] != 1:
            self.player_x = new_x
            self.player_y = new_y
            if self.canvas:
                self.draw_player()
            if self.maze[self.player_y][self.player_x] == 2:
                messagebox.showinfo("胜利", f"你成功走出迷宫 {self.maze_index + 1}！")
                width, height = self.get_maze_size()
                self.generate_single_maze(width, height)

    def toggle_ai(self):
        self.ai_enabled = not self.ai_enabled
        for game in games:
            game.ai_enabled = self.ai_enabled
            if game.ai_button:
                if game.ai_enabled:
                    game.ai_button.config(text="关闭神经网络")
                else:
                    game.ai_button.config(text="启用神经网络")
            game.ai_steps = 0
            game.start_time = time.time()
            if game.ai_thread is None or not game.ai_thread.is_alive():
                game.ai_thread = threading.Thread(target=game.run_ai)
                game.ai_thread.daemon = True
                game.ai_thread.start()

    def set_ai_speed(self, speed):
        if self.maze_index == 0:
            self.ai_speed = int(speed)
            for game in games:
                game.ai_speed = self.ai_speed

    def run_ai(self):
        while self.ai_enabled:
            elapsed_time = time.time() - self.start_time
            if elapsed_time > self.time_limit:
                self.ai_score -= 50
                self.score_label.config(text=f"AI 分数 (迷宫 {self.maze_index + 1}): {self.ai_score}")
                width, height = self.get_maze_size()
                self.generate_single_maze(width, height)
                time.sleep(1 / self.ai_speed)
                continue

            state = self.get_state()
            state_tensor = torch.FloatTensor(state).to(device)
            action_probs = self.model(state_tensor)

            # ε - 贪心策略选择动作
            if random.random() < self.epsilon:
                action = random.randint(0, 3)
            else:
                valid_actions = []
                for action in range(4):
                    dx, dy = [(0, -1), (0, 1), (-1, 0), (1, 0)][action]
                    new_x = self.player_x + dx
                    new_y = self.player_y + dy
                    if 0 <= new_x < self.maze.shape[1] and 0 <= new_y < self.maze.shape[0] and self.maze[new_y][new_x] != 1:
                        valid_actions.append(action)

                if not valid_actions:
                    # 没有可行动作，随机选一个
                    action = random.randint(0, 3)
                else:
                    # 优先选择可行路线上的动作
                    if self.shortest_path and self.path_index < len(self.shortest_path) - 1:
                        next_pos = self.shortest_path[self.path_index + 1]
                        for valid_action in valid_actions:
                            dx, dy = [(0, -1), (0, 1), (-1, 0), (1, 0)][valid_action]
                            new_x = self.player_x + dx
                            new_y = self.player_y + dy
                            if (new_x, new_y) == next_pos:
                                action = valid_action
                                break
                        else:
                            # 如果可行路线上的动作不可行，选择其他可行动作
                            for valid_action in valid_actions:
                                dx, dy = [(0, -1), (0, 1), (-1, 0), (1, 0)][valid_action]
                                new_x = self.player_x + dx
                                new_y = self.player_y + dy
                                if (new_x, new_y) not in self.recent_visited:
                                    action = valid_action
                                    break
                            else:
                                action = random.choice(valid_actions)
                    else:
                        action = random.choice(valid_actions)

            dx, dy = [(0, -1), (0, 1), (-1, 0), (1, 0)][action]
            new_x = self.player_x + dx
            new_y = self.player_y + dy
            reward = 0
            if 0 <= new_x < self.maze.shape[1] and 0 <= new_y < self.maze.shape[0] and self.maze[new_y][new_x] != 1:
                # 中间奖励：接近终点给予奖励
                end_x, end_y = self.maze.shape[1] - 2, self.maze.shape[0] - 2
                old_distance = abs(self.player_x - end_x) + abs(self.player_y - end_y)
                new_distance = abs(new_x - end_x) + abs(new_y - end_y)
                if new_distance < old_distance:
                    reward = 1
                elif new_distance > old_distance:
                    reward = -1  # 远离终点给予惩罚

                # 检查是否沿着最短路径移动
                if self.shortest_path and self.path_index < len(self.shortest_path) - 1:
                    next_pos = self.shortest_path[self.path_index + 1]
                    if (new_x, new_y) == next_pos:
                        reward += 5  # 加大沿着最短路径移动的奖励
                        self.path_index += 1
                    else:
                        reward -= 10  # 加大偏离最短路径的惩罚

                # 探索奖励
                if (new_x, new_y) not in self.visited:
                    reward += 1
                    self.visited.add((new_x, new_y))

                self.player_x = new_x
                self.player_y = new_y
                if self.canvas:
                    def update_canvas():
                        if self.canvas:
                            self.draw_player()
                    self.root.after(0, update_canvas)
                self.recent_visited.append((new_x, new_y))
                self.ai_steps += 1
                if self.ai_steps > self.max_steps:
                    self.ai_score -= 20
                    self.score_label.config(text=f"AI 分数 (迷宫 {self.maze_index + 1}): {self.ai_score}")
                    width, height = self.get_maze_size()
                    self.generate_single_maze(width, height)
                    time.sleep(1 / self.ai_speed)
                    continue
                if self.maze[self.player_y][self.player_x] == 2:
                    reward = max(200, 300 - self.ai_steps)  # 到达终点奖励
                    self.ai_score += reward
                    self.score_label.config(text=f"AI 分数 (迷宫 {self.maze_index + 1}): {self.ai_score}")
                    width, height = self.get_maze_size()
                    self.generate_single_maze(width, height)
                    self.save_model()
                    time.sleep(1 / self.ai_speed)
                    continue
                next_state = self.get_state()
                next_state_tensor = torch.FloatTensor(next_state).to(device)
                target = reward + 0.9 * torch.max(self.model(next_state_tensor)).item()
                loss = nn.MSELoss()(action_probs[action], torch.tensor(target).to(device))
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
                self.save_model()
            time.sleep(1 / self.ai_speed)

    def get_state(self):
        up = self.maze[self.player_y - 1][self.player_x] if self.player_y > 0 else 1
        down = self.maze[self.player_y + 1][self.player_x] if self.player_y < self.maze.shape[0] - 1 else 1
        left = self.maze[self.player_y][self.player_x - 1] if self.player_x > 0 else 1
        right = self.maze[self.player_y][self.player_x + 1] if self.player_x < self.maze.shape[1] - 1 else 1
        return [up, down, left, right]

    def save_model(self):
        torch.save(self.model.state_dict(), "maze_ai_model.pth")

    def load_model(self):
        if os.path.exists("maze_ai_model.pth"):
            self.model.load_state_dict(torch.load("maze_ai_model.pth", map_location=device))

    def get_maze_size(self):
        if self.maze_index == 0:
            try:
                width = int(self.width_entry.get())
                height = int(self.height_entry.get())
                return width, height
            except ValueError:
                # 默认值
                return 10, 10
        else:
            return self.first_game.get_maze_size()


if __name__ == "__main__":
    root = tk.Tk()
    parallel_frame = tk.Frame(root)
    parallel_frame.pack()
    tk.Label(parallel_frame, text="并行迷宫数量:").grid(row=0, column=0)
    parallel_entry = tk.Entry(parallel_frame)
    parallel_entry.grid(row=0, column=1)
    tk.Button(parallel_frame, text="确定", command=lambda: create_mazes(parallel_entry.get())).grid(row=0, column=2)

    show_all = tk.BooleanVar()
    show_all.set(False)
    checkbutton = tk.Checkbutton(root, text="显示所有迷宫", variable=show_all)
    checkbutton.pack()

    games = []


    def create_mazes(num_str):
        try:
            num_mazes = int(num_str)
            for game in games:
                if game.canvas:
                    game.canvas.destroy()
                if game.input_frame:
                    game.input_frame.destroy()
                if game.control_frame:
                    game.control_frame.destroy()
                game.score_label.destroy()
            games.clear()
            first_game = None
            for i in range(num_mazes):
                if i == 0:
                    game = MazeGame(root, i, num_mazes, show_all.get())
                    first_game = game
                else:
                    game = MazeGame(root, i, num_mazes, show_all.get(), first_game)
                games.append(game)
        except ValueError:
            messagebox.showerror("输入错误", "请输入有效的整数并行迷宫数量")


    show_all.trace_add("write", lambda *args: update_mazes())


    def update_mazes():
        for game in games:
            game.show_all = show_all.get()
            if game.show_all or game.maze_index == 0:
                if not game.canvas:
                    game.canvas = tk.Canvas(root, width=300, height=300)
                    game.canvas.pack(side=tk.LEFT, padx=5, pady=5)
                    width, height = game.get_maze_size()
                    game.generate_single_maze(width, height)
            else:
                if game.canvas:
                    game.canvas.destroy()
                    game.canvas = None


    root.mainloop()